import torch
from torch import nn
from torchvision.models import resnet50

from network.backbone.bottleneck_transformer_pytorch import BottleStack

layer = BottleStack(
    dim = 256,
    fmap_size = (129,129),        # set specifically for imagenet's 224 x 224
    dim_out = 2048,
    proj_factor = 4,
    downsample = True,
    heads = 4,
    dim_head = 128,
    rel_pos_emb = True,
    activation = nn.ReLU()
)

resnet = resnet50()

# model surgery

backbone = list(resnet.children())

model = nn.Sequential(
    *backbone[:4],
    #layer,
)

# use the 'BotNet'

img = torch.randn(2, 3, 129, 129)
preds = model(img) # (2, 1000)
print(preds.shape)